// to compile in linux machines: nvcc -o randWalk randWalk.cu -arch=sm_11 -lm -lpng -O3 
// to execute: ./randWalk arg1 arg2 arg3
// arg1: matrix size
// arg2: particle number (power of 2) e.g. 2⁶, 2^7, 2^8... 
// arg3: output image name

#include<stdio.h>
#include<stdlib.h>
#include<cuda.h>
#include<curand_kernel.h>
#include <time.h>
#include <malloc.h>
#include <png.h>

#define CUDA_CALL(lin) do { if((lin) != cudaSuccess) { \
printf("Error at %s:%d\n",__FILE__,__LINE__);\
return EXIT_FAILURE;}} while(0)


#define HANDLE_ERROR( err ) (HandleError( err, __FILE__, __LINE__ ))

inline void setRGB (png_byte *ptr, float val);
int writeImage(char* filename, int *buffer, char* title, int Size);

static void HandleError( cudaError_t err,
                         const char *file,
                         int line ) {
    if (err != cudaSuccess) {
        printf( "%s in %s at line %d\n", cudaGetErrorString( err ),
                file, line );
        exit( EXIT_FAILURE );
    }
}

__global__ void setup_kernel ( curandState* state, unsigned long long seed )
{
    int id = blockIdx.x * blockDim.x + threadIdx.x;
    
    /* Each thread gets same seed, a different sequence number , no offset */
    curand_init ( seed, id, 0, &state[id] ) ;
}

__global__ void randomWalkGpu( int* matriz, const int size, curandState* state )
{
    int id = blockIdx.x * blockDim.x + threadIdx.x;
    int ok;
    int passo = 0;
    int maxSteps = size*size;
    int matriz_id;
    
    /* Cria variáveis para controlar a movimentação da particula nos registradores
    de modo que cada thread tenha sua própria cópia das variáveis que tem acesso
    super rápido */
    int lin = 0;
    int col = 0;

    unsigned int random_number;
    
    /* Copy state to local memory for efficiency */
    curandState localState = state [ id ];
    
	atomicAdd ( &matriz[0], 1 );
	
    while( passo < maxSteps )
    {
        random_number = curand(&localState) % 4;

        ok = 0;

        switch ( random_number ){
            case 0:
                if ( ( lin - 1 ) >= 0 ){
                    --lin;
                    ok = 1;
                }
                break;
            case 1:
                if ( ( col + 1 ) < size ){
                    ++col;
                    ok = 1;
                }
                break;
            case 2:
                if ( ( lin + 1) < size ){
                    ++lin;
                    ok = 1;
                }
                break;
            case 3:
                if ( ( col - 1 ) >= 0 ){
                    --col;
                    ok = 1;
                }
                break;
        }
        
        if ( ok ){
            passo++;
            matriz_id = (lin * size) + col;

            // operação atomic
            atomicAdd ( &matriz[matriz_id], 1) ;
            
            if ( matriz_id == (maxSteps - 1) )
                break;
        }

    }
}

/*
for( passo = 0; passo < maxSteps; )
{
    random_number = curand_uniform(&localState);

    ok = 0;
    
    if (random_number < .25 ){
        if ( ( lin - 1 ) >= 0 ){
            --lin;
            ok = 1;
        }
    }
    else
        if (random_number > .25 && random_number < .5 ){
            if ( ( col + 1 ) < size ){
                ++col;
                ok = 1;
            }
        }
        else
            if (random_number > .5 && random_number < .75 ){
                if ( ( lin + 1) < size ){
                    ++lin;
                    ok = 1;
                }
            
            }
            else{
                if ( ( col - 1 ) >= 0 ){
                    --col;
                    ok = 1;
                }
            }
    */


void print_malha( int* matriz, int size )
{
    int i, j;
    for ( i = 0; i < size; i++ ){
        for ( j = 0; j < size; j++ ){
                printf("%5d", matriz[j + i * size] );

        }
        printf("\n");
    }
}


int main( int argc, char** argv )
{      
    // ./<executable_name> <tamanho_da_matriz> <número_de_particulas> <nome_do_arquivo_de_saída>
    if ( argc < 4 ){
        printf("Entre com os argumentos necessários.\n");
        printf("./<executable_name> <tamanho_da_matriz> <número_de_blocos> <numero_de_threads_por_bloco> <nome_do_arquivo_de_saída>\n");
        return EXIT_FAILURE;
    }
    
    
    int size = atoi(argv[1]); // recupera o tamanho da matriz
    int blockNum = atoi(argv[2]); // recupera o número de blocos
    int threadsNum = atoi(argv[3]); // recupera o número de threads por bloco
    char outputName[50] = "arquivo_saida.png";   /* Nome de saida padrão para a imagem a gerada */
    strcpy(outputName, argv[4]); // captura o nome do arquivo de saída
    int totalParticulas = blockNum * threadsNum; // o total de particulas é equivalente ao total de threads disparados
    
    printf("- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n");
    printf("RW-Gpu\nSize: %d x %d\nParticulas: %d\n\n", size, size, totalParticulas );
            
            
    /*Cria uma matriz no host para receber os valores cauculados no dispositivo*/
    int* buffer_h;
    buffer_h = (int*) calloc ( size * size, sizeof (int) );
    
    // aloca memoria no despositivo cuda 
    int *devMatrix;
    HANDLE_ERROR(cudaMalloc((void**)&devMatrix, size * size * sizeof(int) ) );
    HANDLE_ERROR(cudaMemset( (void*)devMatrix, 0, size * size * sizeof(int) ) );//inicializa a matrix com 0
    
    curandState * devStates;
    HANDLE_ERROR(cudaMalloc((void **)&devStates, totalParticulas * sizeof(curandState)));
    
                        
    setup_kernel<<<blockNum, threadsNum>>>
        ( devStates, time(NULL) );
        
        
    // capture the start time
    // starting the timer here so that we include the cost of
    // all of the operations on the GPU.  if the data were
    // already on the GPU and we just timed the kernel
    // the timing would drop from 74 ms to 15 ms.  Very fast.
    cudaEvent_t     start, stop;
    HANDLE_ERROR( cudaEventCreate( &start ) );
    HANDLE_ERROR( cudaEventCreate( &stop ) );
    HANDLE_ERROR( cudaEventRecord( start, 0 ) ); 
    
                         
    randomWalkGpu <<<blockNum, threadsNum>>>
          (devMatrix, size, devStates);
    const char * error;
    error = cudaGetErrorString(cudaPeekAtLastError());
    printf("%s\n", error);
                    
                                
    // get stop time, and display the timing results
    HANDLE_ERROR( cudaEventRecord( stop, 0 ) );
    HANDLE_ERROR( cudaEventSynchronize( stop ) );
    float   elapsedTime;
    HANDLE_ERROR( cudaEventElapsedTime( &elapsedTime,
                                        start, stop ) );

    printf("Tempo Gpu (em Segundos) = %.3fs\n", elapsedTime/1000.0f);
    
    
    /*Cópia dados da matriz do device para o dispositivo*/
    HANDLE_ERROR( cudaMemcpy ( buffer_h, devMatrix, size * size * sizeof(int), cudaMemcpyDeviceToHost ) );
                                    
    printf("Chegaram %d Particulas\n",  buffer_h[size * size - 1] );
    printf("Probabilidade: %lf\n", (buffer_h[size*size-1] / (double)totalParticulas) );
    //print_malha( buffer_h, size );
    
    int result = writeImage( outputName, buffer_h, "This is a Random Walk result's image", size);

    cudaFree( devMatrix );
    free(buffer_h);
    
	return EXIT_SUCCESS;
}


int writeImage(char* filename, int *buffer, char* title, int Size)
{
    int x, y;
    int maior = buffer[0];
    int menor = buffer[0];
    int code = 0;
    FILE *fp;
    png_structp png_ptr;
    png_infop info_ptr;
    png_bytep row;

    // Open file for writing (binary mode)
    fp = fopen(filename, "wb");
    if (fp == NULL) {
	    fprintf(stderr, "Could not open file %s for writing\n", filename);
	    code = 1;
	    goto finalise;
    }

    // Initialize write structure
    png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
    if (png_ptr == NULL) {
	    fprintf(stderr, "Could not allocate write struct\n");
	    code = 1;
	    goto finalise;
    }

    // Initialize info structure
    info_ptr = png_create_info_struct(png_ptr);
    if (info_ptr == NULL) {
	    fprintf(stderr, "Could not allocate info struct\n");
	    code = 1;
	    goto finalise;
    }

    // Setup Exception handling
    if (setjmp(png_jmpbuf(png_ptr))) {
	    fprintf(stderr, "Error during png creation\n");
	    code = 1;
	    goto finalise;
    }

    png_init_io(png_ptr, fp);

    // Write header (8 bit colour depth)
    png_set_IHDR(png_ptr, info_ptr, Size, Size,
		    8, PNG_COLOR_TYPE_RGB, PNG_INTERLACE_NONE,
		    PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);

    // Set title
    if (title != NULL) {
	    png_text title_text;
	    title_text.compression = PNG_TEXT_COMPRESSION_NONE;
	    title_text.key = "Title";
	    title_text.text = title;
	    png_set_text(png_ptr, info_ptr, &title_text, 1);
    }

    png_write_info(png_ptr, info_ptr);

    // Allocate memory for one row (3 bytes per pixel - RGB)
    row = (png_bytep) malloc(3 * Size * sizeof(png_byte));

    //define maio e menor elemento
    //********************************
    int i;
    for ( i = 0; i < (Size*Size); i++ )
	if ( buffer[i] > maior )
	    maior = buffer[i];
	    else
		if ( buffer[i] < menor )
		    menor = buffer[i];
    //********************************

    float num;
    for (y=0 ; y<Size ; y++) {
	    for (x=0 ; x<Size ; x++) {
		    num = (float)((float)( buffer[y*Size + x] - menor)/(maior - menor) );
		    setRGB(&(row[x*3]), num );
	    }
	    png_write_row(png_ptr, row);
    }

    // End write
    png_write_end(png_ptr, NULL);

    finalise:
    if (fp != NULL) fclose(fp);
    if (info_ptr != NULL) png_free_data(png_ptr, info_ptr, PNG_FREE_ALL, -1);
    if (png_ptr != NULL) png_destroy_write_struct(&png_ptr, (png_infopp)NULL);
    if (row != NULL) free(row);

    return code;
}

inline void setRGB (png_byte *ptr, float val)
{
	int offset;
	double v = (double)(val * 768);
	if (v > 768) v = 768;

	if ( v != 0 )
	{
	    //printf("Valor Diferente de Zero: %f\n", v );
	    int valor = (int) v;

	    if (valor < 1) valor = 1;
	    offset = valor % 256;

	    if (valor < 256) {
		    ptr[0] = 0; ptr[1] = 0; ptr[2] = offset;
	    }
	    else if (valor < 512) {
		    ptr[0] = 0; ptr[1] = offset; ptr[2] = 255-offset;
	    }
	    else {
		    ptr[0] = offset; ptr[1] = 255-offset; ptr[2] = 0;
	    }
	}
	else
	{
	    v = 0;

	    if ( v == 0 ){
		    ptr[0] = 255; ptr[1] = 255; ptr[2] = 255;
		    //cont++;
	    }
	}
}
